package com.leetroll.media.remove;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.imageio.ImageIO;
import javax.imageio.ImageWriter;
import javax.imageio.stream.ImageOutputStream;
import javax.net.ssl.HttpsURLConnection;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.net.HttpURLConnection;
import java.net.URL;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

public class WatermarkRemove {

    public static final Logger logger = LoggerFactory.getLogger(WatermarkRemove.class);

    private static List<File> fileList = new ArrayList<>();

    public static void main(String[] args) {
        convertAllImages("C:\\Users\\86185\\Desktop\\leetroll\\素材\\图片-需抠图",
                "C:\\Users\\86185\\Desktop\\leetroll\\素材\\finish"); // 支持批量去除图片水印
//        convertAllImages("D:\\fan.zhou\\可行性研究\\正时\\Tmp-图片去水印\\timing_pic\\",
//                "D:\\fan.zhou\\可行性研究\\正时\\Tmp-图片去水印\\timing_pic_handle\\"); // 支持批量去除图片水印
    }

    private static void convertAllImages(String dir, String saveDir) {
        File dirFile = new File(dir);
        dir = dirFile.getAbsolutePath();
        File saveDirFile = new File(saveDir);
        saveDir = saveDirFile.getAbsolutePath();
        loadImages(new File(dir));

        for (File file : fileList) {
            String filePath = file.getAbsolutePath();
            String dstPath = saveDir + filePath.substring(filePath.indexOf(dir) + dir.length(), filePath.length());
            logger.info("Converting: {}", filePath);
            replaceColor(file.getAbsolutePath(), dstPath);
        }
    }

    public static void loadImages(File f) {
        if (f != null) {
            if (f.isDirectory()) {
                File[] fileArray = f.listFiles();
                if (fileArray != null) {
                    for (int i = 0; i < fileArray.length; i++) {
                        loadImages(fileArray[i]);
                    }
                }
            } else {
                String name = f.getName();
                if (name.endsWith("png") || name.endsWith("jpg")) {
                    fileList.add(f);
                }
            }
        }
    }

    private static void replaceColor(String srcFile, String dstFile) {
        try {
            replaceImageColor(srcFile, dstFile);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public static void replaceImageColor(String file, String dstFile) throws IOException {
        URL http;
        if (file.trim().startsWith("https")) {
            http = new URL(file);
            HttpsURLConnection conn = (HttpsURLConnection) http.openConnection();
            conn.setRequestMethod("GET");
        } else if (file.trim().startsWith("http")) {
            http = new URL(file);
            HttpURLConnection conn = (HttpURLConnection) http.openConnection();
            conn.setRequestMethod("GET");
        } else {
            http = new File(file).toURI().toURL();
        }
        BufferedImage bi = ImageIO.read(http.openStream());
        if (bi == null) {
            return;
        }

        int pointNum = 30;
        double[][] xy = new double[2][4 * pointNum];
        double[][] rgb = new double[3][4 * pointNum];
        int widthBorder = 30;
        int width = bi.getWidth();
        double widthStep = (double) (width - widthBorder * 2) / (pointNum - 1);
        int heightBorder = 30;
        int height = bi.getHeight();
        double heightStep = (double) (height - heightBorder * 2) / (pointNum - 1);
        for (int i = 0; i < pointNum; i++) {
            xy[0][i] = widthBorder;
            xy[1][i] = (int) (heightBorder + heightStep * i);
            Color oriColor = new Color(bi.getRGB((int) xy[0][i], (int) xy[1][i]));
            rgb[0][i] = oriColor.getRed();
            rgb[1][i] = oriColor.getGreen();
            rgb[2][i] = oriColor.getBlue();

            xy[0][pointNum + i] = width - widthBorder;
            xy[1][pointNum + i] = (int) (heightBorder + heightStep * i);
            oriColor = new Color(bi.getRGB((int) xy[0][pointNum + i], (int) xy[1][pointNum + i]));
            rgb[0][pointNum + i] = oriColor.getRed();
            rgb[1][pointNum + i] = oriColor.getGreen();
            rgb[2][pointNum + i] = oriColor.getBlue();
        }
        for (int i = 0; i < pointNum; i++) {
            xy[0][2 * pointNum + i] = (int) (widthBorder + widthStep * i);
            xy[1][2 * pointNum + i] = heightBorder;
            Color oriColor = new Color(bi.getRGB((int) xy[0][2 * pointNum + i], (int) xy[1][2 * pointNum + i]));
            rgb[0][2 * pointNum + i] = oriColor.getRed();
            rgb[1][2 * pointNum + i] = oriColor.getGreen();
            rgb[2][2 * pointNum + i] = oriColor.getBlue();

            xy[0][3 * pointNum + i] = (int) (widthBorder + widthStep * i);
            xy[1][3 * pointNum + i] = height - heightBorder;
            oriColor = new Color(bi.getRGB((int) xy[0][3 * pointNum + i], (int) xy[1][3 * pointNum + i]));
            rgb[0][3 * pointNum + i] = oriColor.getRed();
            rgb[1][3 * pointNum + i] = oriColor.getGreen();
            rgb[2][3 * pointNum + i] = oriColor.getBlue();
        }
        double[] ar = LinearRegression.getRegressionCoefficient(xy, rgb[0]);
        double[] ag = LinearRegression.getRegressionCoefficient(xy, rgb[1]);
        double[] ab = LinearRegression.getRegressionCoefficient(xy, rgb[2]);

        int colorDiffInit = 20;
        for (int i = 0; i < bi.getWidth(); i++) {
            for (int j = 0; j < bi.getHeight(); j++) {
                int redStd = (int) (i * ar[0] + j * ar[1] + ar[2]);
                redStd = redStd <= 255 ? redStd : 255;
                int greedStd = (int) (i * ag[0] + j * ag[1] + ag[2]);
                greedStd = greedStd <= 255 ? greedStd : 255;
                int blueStd = (int) (i * ab[0] + j * ab[1] + ab[2]);
                blueStd = blueStd <= 255 ? blueStd : 255;

                Color oriColor = new Color(bi.getRGB(i, j));
                int red = oriColor.getRed();
                int greed = oriColor.getGreen();
                int blue = oriColor.getBlue();

                int colorDiffRed = redStd > red ? redStd - red : red - redStd;
                int colorDiffGreed = greedStd > greed ? greedStd - greed : greed - greedStd;
                int colorDiffBlue = blueStd > blue ? blueStd - blue : blue - blueStd;
                if (colorDiffRed < colorDiffInit
                        && colorDiffGreed < colorDiffInit
                        && colorDiffBlue < colorDiffInit) {
                    Color wColor;
                    try {
                        wColor = new Color(redStd, greedStd, blueStd);
                    } catch (Exception ex) {
                        logger.error("", ex);
                        continue;
                    }
                    bi.setRGB(i, j, wColor.getRGB());
                }
            }
        }

        String type = file.substring(file.lastIndexOf(".") + 1, file.length());
        Iterator<ImageWriter> it = ImageIO.getImageWritersByFormatName(type);
        ImageWriter writer = it.next();
        File f = new File(dstFile);
        f.getParentFile().mkdirs();
        ImageOutputStream ios = ImageIO.createImageOutputStream(f);
        writer.setOutput(ios);
        writer.write(bi);
        bi.flush();
        ios.flush();
        ios.close();
    }

}

